package com.agentframework.core.provider;

import com.agentframework.core.agent.AbstractAgent;
import com.agentframework.core.agent.Agent;
import com.agentframework.core.agent.AgentHolder;
import com.agentframework.core.exception.AgentException;
import com.agentframework.core.memory.AgentInMemoryChatMemory;
import io.github.resilience4j.circuitbreaker.CircuitBreaker;
import io.github.resilience4j.circuitbreaker.CircuitBreakerConfig;
import jakarta.annotation.PostConstruct;
import jakarta.annotation.PreDestroy;
import org.jetbrains.annotations.NotNull;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.DefaultChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.api.BaseAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.PromptTemplate;
import org.springframework.ai.chat.prompt.SystemPromptTemplate;
import org.springframework.ai.content.Media;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.util.StringUtils;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.core.scheduler.Schedulers;
import reactor.util.retry.Retry;

import java.io.IOException;
import java.time.Duration;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Consumer;
import java.util.function.Function;

/**
 * @author wangjun
 **/
public class Llm {

    private static final Logger logger = LoggerFactory.getLogger(Llm.class);

    private final LlmProperties llmProperties;

    private final Provider provider;

    private final ChatMemory chatMemory;

    private volatile LlmType llmType;

    private volatile CircuitBreaker circuitBreaker;

    // 线程池用于隔离执行LLM调用
    private final ExecutorService executor = new ThreadPoolExecutor(
            2, 5,
            60L, TimeUnit.SECONDS,
            new LinkedBlockingQueue<>(20),
            new ThreadPoolExecutor.CallerRunsPolicy());

    public Llm(Provider provider, ChatMemory chatMemory, LlmProperties llmProperties) {
        this.provider = provider;
        this.chatMemory = chatMemory;
        this.llmProperties = llmProperties;
    }

    @PostConstruct
    public void init() {
        if (llmProperties != null && llmProperties.getCircuitBreaker().enabled()) {
            CircuitBreakerConfig circuitBreakerConfig = CircuitBreakerConfig.custom()
                    .failureRateThreshold(llmProperties.getCircuitBreaker().failureRate())
                    .slidingWindowType(CircuitBreakerConfig.SlidingWindowType.COUNT_BASED)
                    .slidingWindowSize(10)
                    .minimumNumberOfCalls(5)
                    .waitDurationInOpenState(Duration.ofSeconds(llmProperties.getCircuitBreaker().waitDurationSeconds()))
                    .permittedNumberOfCallsInHalfOpenState(2)
                    .recordExceptions(
                            TimeoutException.class,
                            IOException.class,
                            RuntimeException.class)
                    .build();

            circuitBreaker = CircuitBreaker.of("llmCircuitBreaker", circuitBreakerConfig);
        }
    }

    public Message call(Prompt prompt, Consumer<ChatClient.AdvisorSpec> advisorSpecConsumer, List<ToolCallback> toolCallbacks, Map<String, Object> toolContext) {

        try {
            // 准备客户端 TODO 此处后续再考虑做缓存
            ChatClient chatClient = prepareClient(advisorSpecConsumer, toolCallbacks, toolContext);

            // 执行调用
            Function<Prompt, Mono<AssistantMessage>> decoratedCall = p -> executeWithRetry(chatClient, p);

            // 如果启用了断路器，应用断路器保护
            if (llmProperties.getCircuitBreaker().enabled() && circuitBreaker != null) {
                decoratedCall = CircuitBreaker.decorateFunction(circuitBreaker, decoratedCall);
            }

            // 执行调用并处理结果
            return decoratedCall.apply(prompt)
                    .timeout(Duration.ofSeconds(llmProperties.getTimeout().seconds()))
                    .onErrorResume(TimeoutException.class, e -> {
                        logger.error("LLM call timed out after {} seconds", llmProperties.getTimeout().seconds(), e);
                        return Mono.just(new AssistantMessage("很抱歉，系统响应超时，请稍后重试"));
                    })
                    .onErrorResume(Exception.class, e -> {
                        logger.error("LLM call failed", e);
                        return Mono.just(new AssistantMessage("系统暂时无法响应，请稍后重试"));
                    })
                    .block();
        } catch (Exception e) {
            logger.error("Unexpected error in LLM call", e);
            throw new AgentException(e);
        }

    }

    private ChatClient prepareClient(Consumer<ChatClient.AdvisorSpec> advisorSpecConsumer, List<ToolCallback> toolCallbacks, Map<String, Object> toolContext) {
        var advisorSpec = new DefaultChatClient.DefaultAdvisorSpec();
        advisorSpecConsumer.accept(advisorSpec);

        ChatClient.Builder builder = ChatClient.builder(provider.getChatModel(getLlmType()));
        String systemPrompt = systemPrompt();
        if (StringUtils.hasText(systemPrompt)) {
            builder.defaultSystem(systemPrompt);
        }

        return builder
                .defaultAdvisors(MessageChatMemoryAdvisor.builder(((AgentInMemoryChatMemory) chatMemory).setAgent(getAgent())).conversationId(getEnvId()).scheduler(BaseAdvisor.DEFAULT_SCHEDULER).build())
                .defaultAdvisors(advisorSpec.getAdvisors())
                .defaultToolCallbacks(toolCallbacks)
                .defaultToolCallbacks(getToolCallbacks())
                .defaultToolContext(toolContext)
                .defaultToolContext(getToolContext())
                .build();
    }

    /**
     * 带重试的LLM调用执行
     */
    private Mono<AssistantMessage> executeWithRetry(ChatClient chatClient, Prompt prompt) {
        return Mono.defer(() -> {
                    logger.debug("{} Starting LLM call", getAgent().toString());

                    // 把实际的调用放到另一个线程执行，避免阻塞
                    return Mono.fromCallable(() -> {
                                try {
                                    Flux<String> contentFlux = chatClient.prompt(prompt).stream().content()
                                            .onBackpressureBuffer(llmProperties.getStream().bufferSize());

                                    // 收集并处理响应
                                    List<String> contentList = new ArrayList<>();
                                    CountDownLatch latch = new CountDownLatch(1);
                                    AtomicReference<Throwable> error = new AtomicReference<>();

                                    contentFlux.subscribe(
                                            // 处理每一段内容
                                            content -> {
                                                if (StringUtils.hasText(content)) {
                                                    System.out.print(content);
                                                    contentList.add(content);
                                                }
                                            },
                                            // 处理错误
                                            err -> {
                                                error.set(err);
                                                latch.countDown();
                                            },
                                            // 处理完成
                                            () -> {
                                                logger.debug("[Response Complete]");
                                                latch.countDown();
                                            }
                                    );

                                    // 等待响应完成或超时
                                    if (!latch.await(llmProperties.getTimeout().seconds(), TimeUnit.SECONDS)) {
                                        throw new TimeoutException("LLM response timed out");
                                    }

                                    // 检查是否有错误
                                    if (error.get() != null) {
                                        throw new RuntimeException("Error in LLM call", error.get());
                                    }

                                    // 处理响应
                                    String content = String.join("", contentList);
                                    return new AssistantMessage(content);
                                } catch (Exception e) {
                                    throw new RuntimeException("Error in LLM call", e);
                                }
                            })
                            .subscribeOn(Schedulers.fromExecutor(executor));
                })
                .retryWhen(Retry.backoff(llmProperties.getRetry().max(), Duration.ofSeconds(llmProperties.getRetry().initialDelaySeconds()))
                        .filter(this::isRetryableError)
                        .doBeforeRetry(retrySignal ->
                                logger.warn("Retrying LLM call after error: {} (attempt {}/{})",
                                        retrySignal.failure().getMessage(),
                                        retrySignal.totalRetries() + 1,
                                        llmProperties.getRetry().max())
                        )
                );
    }

    /**
     * 判断是否可以重试的错误
     */
    private boolean isRetryableError(Throwable throwable) {
        if (throwable instanceof TimeoutException) {
            return true;
        }

        if (throwable instanceof IOException) {
            return true;
        }

        if (throwable instanceof RuntimeException) {
            String message = throwable.getMessage();
            // 排除非重试错误
            return message == null || (!message.contains("Invalid request") &&
                    !message.contains("model not found") &&
                    !message.contains("invalid api key") &&
                    !message.contains("context length"));
        }

        return false;
    }

    public Message call(Prompt prompt, Consumer<ChatClient.AdvisorSpec> advisorSpecConsumer) {
        return call(prompt, advisorSpecConsumer, List.of(), Map.of());
    }

    public Message call(Prompt prompt, List<ToolCallback> toolCallbacks, Map<String, Object> toolContext) {
        return call(prompt, a -> {}, toolCallbacks, toolContext);
    }

    public Message call(Prompt prompt, List<ToolCallback> toolCallbacks) {
        return call(prompt, toolCallbacks, Map.of());
    }

    public Message call(Prompt prompt) {
        return call(prompt, List.of());
    }

    public Message call(Consumer<ChatClient.PromptSystemSpec> systemSpecConsumer, Consumer<ChatClient.PromptUserSpec> userSpecConsumer,
                        Consumer<ChatClient.AdvisorSpec> advisorSpecConsumer, List<ToolCallback> toolCallbacks, Map<String, Object> toolContext) {

        var us = new AgentDefaultPromptUserSpec();
        userSpecConsumer.accept(us);

        PromptTemplate promptTemplate = new PromptTemplate(us.text());
        Message userMessage = promptTemplate.createMessage(us.params());
        if (!us.media().isEmpty()) {
            userMessage = UserMessage.builder().text(userMessage.getText()).media(us.media()).build();
        }

        var systemSpec = new AgentDefaultPromptSystemSpec();
        systemSpecConsumer.accept(systemSpec);

        return call(new Prompt(StringUtils.hasText(systemSpec.text()) ? List.of(new SystemPromptTemplate(systemSpec.text()).createMessage(systemSpec.params()), userMessage) : List.of(userMessage)),
                advisorSpecConsumer, toolCallbacks, toolContext);
    }

    public Message call(Consumer<ChatClient.PromptSystemSpec> systemSpecConsumer, Consumer<ChatClient.PromptUserSpec> userSpecConsumer, Consumer<ChatClient.AdvisorSpec> advisorSpecConsumer) {
        return call(systemSpecConsumer, userSpecConsumer, advisorSpecConsumer, List.of(), Map.of());
    }

    public Message call(Consumer<ChatClient.PromptSystemSpec> systemSpecConsumer, Consumer<ChatClient.PromptUserSpec> userSpecConsumer, List<ToolCallback> toolCallbacks, Map<String, Object> toolContext) {
        return call(systemSpecConsumer, userSpecConsumer, a -> {}, toolCallbacks, toolContext);
    }

    public Message call(Consumer<ChatClient.PromptSystemSpec> systemSpecConsumer, Consumer<ChatClient.PromptUserSpec> userSpecConsumer, List<ToolCallback> toolCallbacks) {
        return call(systemSpecConsumer, userSpecConsumer, toolCallbacks, Map.of());
    }

    public Message call(Consumer<ChatClient.PromptUserSpec> userSpecConsumer) {
        return call(s -> {}, userSpecConsumer, List.of());
    }

    public Message call(Consumer<ChatClient.PromptSystemSpec> systemSpecConsumer, Consumer<ChatClient.PromptUserSpec> userSpecConsumer) {
        return call(systemSpecConsumer, userSpecConsumer, List.of());
    }

    public Message call(String message) {
        return call(PromptTemplate.builder().template(message).variables(new HashMap<>()).renderer((template, variables) -> template).build().create());
    }

    public ChatMemory getChatMemory() {
        return chatMemory;
    }

    public void setLlmType(LlmType llmType) {
        this.llmType = llmType;
    }

    public LlmType getLlmType() {
        if (this.llmType != null) {
            return this.llmType;
        }
        Agent agent = AgentHolder.getAgent();
        return agent == null || ((AbstractAgent) agent).getLlmType() == null ? LlmType.DEEPSEEK : ((AbstractAgent) agent).getLlmType();
    }

    public AbstractAgent getAgent() {
        return (AbstractAgent) AgentHolder.getAgent();
    }

    public String getEnvId() {
        AbstractAgent agent = (AbstractAgent) AgentHolder.getAgent();
        return agent == null ? "" : agent.getEnvId();
    }

    public List<ToolCallback> getToolCallbacks() {
        AbstractAgent agent = (AbstractAgent) AgentHolder.getAgent();
        return agent == null ? List.of() : agent.getToolCallbacks();
    }

    public Map<String, Object> getToolContext() {
        AbstractAgent agent = (AbstractAgent) AgentHolder.getAgent();
        return agent == null ? Map.of() : agent.getToolContext();
    }

    public String systemPrompt() {
        if (AgentHolder.getAgent() instanceof AbstractAgent aagent) {
            Map<String, Object> variablesMap = Map.of(aagent.NAME, aagent.getAgentName(), aagent.PROFILE, aagent.getAgentProfile(), aagent.GOAL, aagent.getGoal());
            return PromptTemplate.builder().template(aagent.PREFIX_TEMPLATE).variables(variablesMap).build().createMessage().getText();
        }
        return null;
    }

    @PreDestroy
    public void destroy() {
        executor.shutdown();
        try {
            if (!executor.awaitTermination(5, TimeUnit.SECONDS)) {
                executor.shutdownNow();
            }
        } catch (InterruptedException e) {
            executor.shutdownNow();
            Thread.currentThread().interrupt();
        }
    }

    static class AgentDefaultPromptSystemSpec extends DefaultChatClient.DefaultPromptSystemSpec {
        @NotNull
        @Override
        protected String text() {
            return super.text();
        }

        @NotNull
        @Override
        protected Map<String, Object> params() {
            return super.params();
        }
    }

    static class AgentDefaultPromptUserSpec extends DefaultChatClient.DefaultPromptUserSpec {
        @NotNull
        @Override
        protected String text() {
            return super.text();
        }

        @NotNull
        @Override
        protected List<Media> media() {
            return super.media();
        }

        @NotNull
        @Override
        protected Map<String, Object> params() {
            return super.params();
        }
    }

}
